{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from tensorboardX import SummaryWriter\n",
    "import tqdm\n",
    "import torch.nn as nn\n",
    "from torch.nn import functional as F\n",
    "import math\n",
    "import torch.backends.cudnn as cudnn\n",
    "from torch.nn import init"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DoubleConv(nn.Module):\n",
    "    \"\"\"(convolution => [BN] => ReLU) * 2\"\"\"\n",
    "\n",
    "    def __init__(self, in_channels, out_channels, mid_channels=None):\n",
    "        super().__init__()\n",
    "        if not mid_channels:\n",
    "            mid_channels = out_channels\n",
    "        self.double_conv = nn.Sequential(\n",
    "            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1),\n",
    "            nn.BatchNorm2d(mid_channels),\n",
    "            nn.ReLU(inplace=True),\n",
    "            \n",
    "            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1),\n",
    "            nn.BatchNorm2d(out_channels),\n",
    "            nn.ReLU(inplace=True)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.double_conv(x)\n",
    "\n",
    "#下采样\n",
    "class Down(nn.Module):\n",
    "    \"\"\"Downscaling with maxpool then double conv\"\"\"\n",
    "\n",
    "    def __init__(self, in_channels, out_channels):\n",
    "        super().__init__()\n",
    "        self.maxpool_conv = nn.Sequential(\n",
    "            nn.MaxPool2d(2),\n",
    "            DoubleConv(in_channels, out_channels)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.maxpool_conv(x)\n",
    "\n",
    "#上采样\n",
    "class Up(nn.Module):\n",
    "    \"\"\"Upscaling then double conv\"\"\"\n",
    "\n",
    "    def __init__(self, in_channels, out_channels, bilinear=True):\n",
    "        super().__init__()\n",
    "\n",
    "        # if bilinear, use the normal convolutions to reduce the number of channels\n",
    "        if bilinear:\n",
    "            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)\n",
    "            self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)\n",
    "        else:\n",
    "            self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)\n",
    "            self.conv = DoubleConv(in_channels, out_channels)\n",
    "\n",
    "    def forward(self, x1, x2):\n",
    "        x1 = self.up(x1)\n",
    "        # input is CHW\n",
    "        diffY = x2.size()[2] - x1.size()[2]\n",
    "        diffX = x2.size()[3] - x1.size()[3]\n",
    "\n",
    "        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,\n",
    "                        diffY // 2, diffY - diffY // 2])\n",
    "        # if you have padding issues, see\n",
    "        x = torch.cat([x2, x1], dim=1)\n",
    "        return self.conv(x)\n",
    "\n",
    "#输出再进行一层卷积操作\n",
    "class OutConv(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels):\n",
    "        super(OutConv, self).__init__()\n",
    "        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.conv(x)\n",
    "\n",
    "    \n",
    "#nonLocal算法\n",
    "class NonLocalBlock(nn.Module):\n",
    "    def __init__(self, channel):\n",
    "        super(NonLocalBlock, self).__init__()\n",
    "        self.inter_channel = channel // 2\n",
    "        self.conv_phi = nn.Conv2d(in_channels=channel, out_channels=self.inter_channel, kernel_size=1, stride=1,padding=0, bias=False)\n",
    "        self.conv_theta = nn.Conv2d(in_channels=channel, out_channels=self.inter_channel, kernel_size=1, stride=1, padding=0, bias=False)\n",
    "        self.conv_g = nn.Conv2d(in_channels=channel, out_channels=self.inter_channel, kernel_size=1, stride=1, padding=0, bias=False)\n",
    "        self.softmax = nn.Softmax(dim=1)\n",
    "        self.conv_mask = nn.Conv2d(in_channels=self.inter_channel, out_channels=channel, kernel_size=1, stride=1, padding=0, bias=False)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # [N, C, H , W]\n",
    "        b, c, h, w = x.size()\n",
    "        # [N, C/2, H * W]\n",
    "        x_phi = self.conv_phi(x).view(b, c, -1)\n",
    "        # [N, H * W, C/2]\n",
    "        x_theta = self.conv_theta(x).view(b, c, -1).permute(0, 2, 1).contiguous()\n",
    "        x_g = self.conv_g(x).view(b, c, -1).permute(0, 2, 1).contiguous()\n",
    "        # [N, H * W, H * W]\n",
    "        mul_theta_phi = torch.matmul(x_theta, x_phi)\n",
    "        mul_theta_phi = self.softmax(mul_theta_phi)\n",
    "        # [N, H * W, C/2]\n",
    "        mul_theta_phi_g = torch.matmul(mul_theta_phi, x_g)\n",
    "        # [N, C/2, H, W]\n",
    "        mul_theta_phi_g = mul_theta_phi_g.permute(0,2,1).contiguous().view(b,self.inter_channel, h, w)\n",
    "        # [N, C, H , W]\n",
    "        mask = self.conv_mask(mul_theta_phi_g)\n",
    "        out = mask + x\n",
    "        return out\n",
    "\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class UNet(nn.Module):\n",
    "    def __init__(self, n_channels = 3, n_classes =2, bilinear=True):\n",
    "        super(UNet, self).__init__()\n",
    "        #属性\n",
    "        self.n_channels = n_channels\n",
    "        self.n_classes = n_classes\n",
    "        self.bilinear = bilinear\n",
    "        #特征提取函数和nonlocal算法\n",
    "        self.inc = DoubleConv(n_channels, 64)\n",
    "        self.down1 = Down(64, 128)\n",
    "        self.down2 = Down(128, 256)\n",
    "        self.down3 = Down(256, 512)\n",
    "        self.nl2 = NonLocalBlock(256) \n",
    "        \n",
    "        factor = 2 if bilinear else 1\n",
    "        self.down4 = Down(512, 1024 // factor)\n",
    "        #上采样和连接函数\n",
    "        self.up1 = Up(1024, 512 // factor, bilinear)\n",
    "        self.up2 = Up(512, 256 // factor, bilinear)\n",
    "        self.up3 = Up(256, 128 // factor, bilinear)\n",
    "        self.up4 = Up(128, 64, bilinear)\n",
    "        \n",
    "        self.outc = OutConv(64, n_classes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        #五个特征值\n",
    "        x1 = self.inc(x)\n",
    "        x2 = self.down1(x1)\n",
    "        x3 = self.down2(x2)\n",
    "        x3 = self.nl2(x3)\n",
    "        x4 = self.down3(x3)\n",
    "        x5 = self.down4(x4)\n",
    "        #连接\n",
    "        x = self.up1(x5, x4)\n",
    "        x = self.up2(x, x3)\n",
    "        x = self.up3(x, x2)\n",
    "        x = self.up4(x, x1)\n",
    "        logits = self.outc(x)\n",
    "        \n",
    "        logits = F.log_softmax(logits, dim=1)\n",
    "        return logits,logits"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 模型训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['KMP_DUPLICATE_LIB_OK']='True'\n",
    "import argparse\n",
    "from torch.utils.data import Dataset\n",
    "from torch.utils.data import DataLoader\n",
    "import glob\n",
    "from torchvision import transforms\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "from imgaug import augmenters as iaa\n",
    "import imgaug as ia\n",
    "import socket\n",
    "from datetime import datetime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def augmentation():\n",
    "    # augment images with spatial transformation: Flip, Affine, Rotation, etc...\n",
    "    # see https://github.com/aleju/imgaug for more details\n",
    "    pass\n",
    "\n",
    "\n",
    "def augmentation_pixel():\n",
    "    # augment images with pixel intensity transformation: GaussianBlur, Multiply, etc...\n",
    "    pass\n",
    "\n",
    "\n",
    "class Data(torch.utils.data.Dataset):\n",
    "    Unlabelled = [0, 0, 0]\n",
    "    sick = [255, 255, 255]\n",
    "    COLOR_DICT = np.array([Unlabelled, sick])\n",
    "\n",
    "    def __init__(self, dataset_path, scale=(320, 320), mode='train'):\n",
    "        super().__init__()\n",
    "        self.mode = mode\n",
    "        self.img_path = dataset_path + '/img'\n",
    "        self.mask_path = dataset_path + '/mask'\n",
    "        self.image_lists, self.label_lists = self.read_list(self.img_path)\n",
    "        self.resize = scale\n",
    "        self.flip = iaa.SomeOf((2, 5), [\n",
    "            iaa.PiecewiseAffine(scale=(0, 0.1), nb_rows=4, nb_cols=4, cval=0),\n",
    "            iaa.Fliplr(0.5),\n",
    "            iaa.Flipud(0.1),\n",
    "            iaa.Affine(rotate=(-20, 20),\n",
    "                       scale={\"x\": (0.8, 1.2), \"y\": (0.8, 1.2)}),\n",
    "            iaa.OneOf([\n",
    "                iaa.GaussianBlur((0, 1.0)),  # blur images with a sigma between 0 and 3.0\n",
    "                iaa.AverageBlur(k=(3, 5)),  # blur image using local means with kernel sizes between 2 and 7\n",
    "                iaa.MedianBlur(k=(3, 5)),  # blur image using local medians with kernel sizes between 2 and 7\n",
    "            ]),\n",
    "            iaa.contrast.LinearContrast((0.5, 1.5))],\n",
    "                               random_order=True)\n",
    "\n",
    "        self.to_tensor = transforms.ToTensor()\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        # load image and crop\n",
    "        img = Image.open(self.image_lists[index]).convert('RGB')\n",
    "        img = img.resize(self.resize)\n",
    "        img = np.array(img)\n",
    "        labels = self.label_lists[index]\n",
    "        #load label\n",
    "        if self.mode != 'test':\n",
    "            label_ori = Image.open(self.label_lists[index]).convert('RGB')\n",
    "            label_ori = label_ori.resize(self.resize)\n",
    "            label_ori = np.array(label_ori)\n",
    "            label = np.ones(shape=(label_ori.shape[0], label_ori.shape[1]), dtype=np.uint8)\n",
    "\n",
    "            #convert RGB  to one hot\n",
    "\n",
    "            for i in range(len(self.COLOR_DICT)):\n",
    "                equality = np.equal(label_ori, self.COLOR_DICT[i])\n",
    "                class_map = np.all(equality, axis=-1)\n",
    "                label[class_map] = i\n",
    "\n",
    "            # augment image and label\n",
    "            if self.mode == 'train':\n",
    "                seq_det = self.flip.to_deterministic()  #固定变换\n",
    "                segmap = ia.SegmentationMapsOnImage(label, shape=label.shape)\n",
    "                img = seq_det.augment_image(img)\n",
    "                label = seq_det.augment_segmentation_maps([segmap])[0].get_arr().astype(np.uint8)\n",
    "\n",
    "            label_img = torch.from_numpy(label.copy()).float()\n",
    "            if self.mode == 'val':\n",
    "                img_num = len(os.listdir(os.path.dirname(labels)))\n",
    "                labels = label_img, img_num\n",
    "            else:\n",
    "                labels = label_img\n",
    "        imgs = img.transpose(2, 0, 1) / 255.0\n",
    "        img = torch.from_numpy(imgs.copy()).float()  #self.to_tensor(img.copy()).float()\n",
    "        return img, labels\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.image_lists)\n",
    "\n",
    "    def read_list(self, image_path):\n",
    "        fold = os.listdir(image_path)\n",
    "        # fold = sorted(os.listdir(image_path), key=lambda x: int(x[-2:]))\n",
    "        # print(fold)\n",
    "\n",
    "        img_list = []\n",
    "        label_list = []\n",
    "        if self.mode == 'train':\n",
    "            for item in fold:\n",
    "                name = item.split('.')[0]\n",
    "                img_list.append(os.path.join(image_path, item))\n",
    "                label_list.append(os.path.join(image_path.replace('img', 'mask'), '{}.png'.format(name)))\n",
    "\n",
    "\n",
    "        elif self.mode == 'val':\n",
    "            for item in fold:\n",
    "                name = item.split('.')[0]\n",
    "                img_list.append(os.path.join(image_path, item))\n",
    "                label_list.append(os.path.join(image_path.replace('img', 'mask'), '{}.png'.format(name)))\n",
    "\n",
    "\n",
    "        elif self.mode == 'test':\n",
    "            for item in fold:\n",
    "                name = item.split('.')[0]\n",
    "                img_list.append(os.path.join(image_path, item))\n",
    "                label_list.append(os.path.join(image_path.replace('img', 'mask'), '{}.png'.format(name)))\n",
    "\n",
    "        return img_list, label_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "ROOT_path = os.getcwd()\n",
    "\n",
    "class DefaultConfig(object):\n",
    "    num_epochs = 30\n",
    "    epoch_start_i = 0\n",
    "    checkpoint_step = 5\n",
    "    validation_step = 1\n",
    "    crop_height = 256\n",
    "    crop_width = 448\n",
    "    batch_size = 2\n",
    "    #训练集所在位置\n",
    "    data = r'train_dataset/'\n",
    "\n",
    "    log_dirs = os.path.join(ROOT_path, 'Log/OCT')\n",
    "\n",
    "    lr = 0.01\n",
    "    lr_mode = 'poly'\n",
    "    net_work = 'Unet'\n",
    "    #net_work= 'MSSeg'  #net_work= 'UNet'\n",
    "\n",
    "    momentum = 0.9  #\n",
    "    weight_decay = 1e-4  #\n",
    "\n",
    "    mode = 'train'\n",
    "    num_classes = 2\n",
    "\n",
    "    k_fold = 4\n",
    "    test_fold = 4\n",
    "    num_workers = 0\n",
    "    #model保存位置\n",
    "    cuda = '0'\n",
    "    use_gpu = True\n",
    "    pretrained_model_path = os.path.join(ROOT_path, 'pretrained', 'resnet34-333f7ec4.pth')\n",
    "    save_model_path = './checkpoints'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.autograd import Variable\n",
    "\n",
    "class Multi_DiceLoss(nn.Module):\n",
    "    def __init__(self, class_num=5, smooth=0.001):\n",
    "        super(Multi_DiceLoss, self).__init__()\n",
    "        self.smooth = smooth\n",
    "        self.class_num = class_num\n",
    "\n",
    "    def forward(self, input, target):\n",
    "        input = torch.exp(input)\n",
    "        Dice = Variable(torch.Tensor([0]).float()).cuda()\n",
    "        for i in range(0, self.class_num):\n",
    "            input_i = input[:, i, :, :]\n",
    "            target_i = (target == i).float()\n",
    "            intersect = (input_i * target_i).sum()\n",
    "            union = torch.sum(input_i) + torch.sum(target_i)\n",
    "            dice = (2 * intersect + self.smooth) / (union + self.smooth)\n",
    "            Dice += dice\n",
    "        dice_loss = 1 - Dice / (self.class_num)\n",
    "        return dice_loss\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import shutil\n",
    "import os.path as osp\n",
    "\n",
    "def save_checkpoint(state, best_pred, epoch, is_best, checkpoint_path, filename='./checkpoint/checkpoint.pth'):\n",
    "    torch.save(state, filename)\n",
    "    if is_best:\n",
    "        shutil.copyfile(filename, osp.join(checkpoint_path, 'best.pth'))\n",
    "    # shutil.copyfile(filename, osp.join(checkpoint_path, 'model_ep{}.pth'.format(epoch+1)))\n",
    "\n",
    "\n",
    "def adjust_learning_rate(opt, optimizer, epoch):\n",
    "    \"\"\"\n",
    "    Sets the learning rate to the initial LR decayed by 10 every 30 epochs(step = 30)\n",
    "    \"\"\"\n",
    "    if opt.lr_mode == 'step':\n",
    "        lr = opt.lr * (0.1 ** (epoch // opt.step))\n",
    "    elif opt.lr_mode == 'poly':\n",
    "        lr = opt.lr * (1 - epoch / opt.num_epochs) ** 0.9\n",
    "    else:\n",
    "        raise ValueError('Unknown lr mode {}'.format(opt.lr_mode))\n",
    "\n",
    "    for param_group in optimizer.param_groups:\n",
    "        param_group['lr'] = lr\n",
    "    return lr\n",
    "\n",
    "\n",
    "def eval_multi_seg(predict, target, num_classes):\n",
    "    # pred_seg=torch.argmax(torch.exp(predict),dim=1).int()\n",
    "    pred_seg = predict.data.cpu().numpy()\n",
    "    label_seg = target.data.cpu().numpy().astype(dtype=np.int64)\n",
    "    assert (pred_seg.shape == label_seg.shape)\n",
    "    acc = (pred_seg == label_seg).sum() / (pred_seg.shape[0] * pred_seg.shape[1] * pred_seg.shape[2])\n",
    "\n",
    "    # Dice = []\n",
    "    # Precsion = []\n",
    "    # Jaccard = []\n",
    "    # Sensitivity=[]\n",
    "    # Specificity=[]\n",
    "\n",
    "    # n = pred_seg.shape[0]\n",
    "    Dice = []\n",
    "    True_label = []\n",
    "    TP = FPN = 0\n",
    "    for classes in range(1, num_classes):\n",
    "        overlap = ((pred_seg == classes) * (label_seg == classes)).sum()\n",
    "        union = (pred_seg == classes).sum() + (label_seg == classes).sum()\n",
    "        Dice.append((2 * overlap + 0.1) / (union + 0.1))\n",
    "        True_label.append((label_seg == classes).sum())\n",
    "        TP += overlap\n",
    "        FPN += union\n",
    "\n",
    "    return Dice, True_label, acc, 2 * TP / (FPN + 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "#评估\n",
    "def val(args, model, dataloader):\n",
    "    print('\\n')\n",
    "    print('Start Validation!')\n",
    "    with torch.no_grad():\n",
    "        model.eval()\n",
    "        tbar = tqdm.tqdm(dataloader, desc='\\r')\n",
    "\n",
    "        total_Dice = []\n",
    "        total_Dice1 = []\n",
    "        total_Dice2 = []\n",
    "        total_Dice3 = []\n",
    "        total_Dice.append(total_Dice1)\n",
    "        total_Dice.append(total_Dice2)\n",
    "        total_Dice.append(total_Dice3)\n",
    "        Acc = []\n",
    "\n",
    "        cur_cube = []\n",
    "        cur_label_cube = []\n",
    "        next_cube = []\n",
    "        counter = 0\n",
    "        end_flag = False\n",
    "\n",
    "        for i, (data, labels) in enumerate(tbar):\n",
    "            # tbar.update()\n",
    "            if torch.cuda.is_available() and args.use_gpu:\n",
    "                data = data.cuda()\n",
    "                label = labels[0].cuda()\n",
    "            slice_num = labels[1][0].long().item()\n",
    "\n",
    "            # get RGB predict image\n",
    "\n",
    "            aux_predict, predicts = model(data)\n",
    "\n",
    "            predict = torch.argmax(torch.exp(predicts), dim=1)\n",
    "            batch_size = predict.size()[0]\n",
    "\n",
    "            counter += batch_size\n",
    "            if counter <= slice_num:\n",
    "                cur_cube.append(predict)\n",
    "                cur_label_cube.append(label)\n",
    "                if counter == slice_num:\n",
    "                    end_flag = True\n",
    "                    counter = 0\n",
    "            else:\n",
    "                last = batch_size - (counter - slice_num)\n",
    "\n",
    "                last_p = predict[0:last]\n",
    "                last_l = label[0:last]\n",
    "\n",
    "                first_p = predict[last:]\n",
    "                first_l = label[last:]\n",
    "\n",
    "                cur_cube.append(last_p)\n",
    "                cur_label_cube.append(last_l)\n",
    "                end_flag = True\n",
    "                counter = counter - slice_num\n",
    "\n",
    "            if end_flag:\n",
    "                end_flag = False\n",
    "                predict_cube = torch.stack(cur_cube, dim=0).squeeze()\n",
    "                label_cube = torch.stack(cur_label_cube, dim=0).squeeze()\n",
    "                cur_cube = []\n",
    "                cur_label_cube = []\n",
    "                if counter != 0:\n",
    "                    cur_cube.append(first_p)\n",
    "                    cur_label_cube.append(first_l)\n",
    "\n",
    "                assert predict_cube.size()[0] == slice_num\n",
    "                Dice, true_label, acc, mean_dice = eval_multi_seg(predict_cube, label_cube, args.num_classes)\n",
    "\n",
    "                for class_id in range(args.num_classes - 1):\n",
    "                    if true_label[class_id] != 0:\n",
    "                        total_Dice[class_id].append(Dice[class_id])\n",
    "                Acc.append(acc)\n",
    "                len0 = len(total_Dice[0]) if len(total_Dice[0]) != 0 else 1\n",
    "                len1 = len(total_Dice[1]) if len(total_Dice[1]) != 0 else 1\n",
    "                len2 = len(total_Dice[2]) if len(total_Dice[2]) != 0 else 1\n",
    "\n",
    "                dice1 = sum(total_Dice[0]) / len0\n",
    "                dice2 = sum(total_Dice[1]) / len1\n",
    "                dice3 = sum(total_Dice[2]) / len2\n",
    "                ACC = sum(Acc) / len(Acc)\n",
    "                tbar.set_description('Mean_D: %3f, Dice1: %.3f, Dice2: %.3f, Dice3: %.3f, ACC: %.3f' % (\n",
    "                    mean_dice, dice1, dice2, dice3, ACC))\n",
    "        print('Mean_Dice:', mean_dice)\n",
    "        print('Dice1:', dice1)\n",
    "        print('Dice2:', dice2)\n",
    "        print('Dice3:', dice3)\n",
    "        print('Acc:', ACC)\n",
    "        return mean_dice, dice1, dice2, dice3, ACC\n",
    "\n",
    "#训练\n",
    "def train(args, model, optimizer, criterion, dataloader_train, dataloader_val, ):\n",
    "    current_time = datetime.now().strftime('%b%d_%H-%M-%S')\n",
    "    #对数目录\n",
    "    log_dir = os.path.join(args.log_dirs, current_time + '_' + socket.gethostname())\n",
    "    writer = SummaryWriter(log_dir=log_dir)\n",
    "    step = 0\n",
    "    best_pred = 0.0\n",
    "    # 对数据进行迭代训练\n",
    "    for epoch in range(args.num_epochs):\n",
    "        #调整学习速率\n",
    "        lr = adjust_learning_rate(args, optimizer, epoch)\n",
    "        #训练\n",
    "        model.train()\n",
    "        tq = tqdm.tqdm(total=len(dataloader_train) * args.batch_size)#进度条\n",
    "        tq.set_description('epoch %d, lr %f' % (epoch, lr))\n",
    "        #损失函数初始化\n",
    "        loss_record = []\n",
    "        train_loss = 0.0\n",
    "        # is_best=False  组成索引序列\n",
    "        for i, (data, label) in enumerate(dataloader_train):\n",
    "            # if i>9:\n",
    "            #     break\n",
    "            if torch.cuda.is_available() and args.use_gpu:\n",
    "                data = data.cuda()\n",
    "                label = label.cuda().long()\n",
    "            optimizer.zero_grad()\n",
    "            aux_out, main_out = model(data)\n",
    "            #损失函数\n",
    "            loss_aux = F.nll_loss(aux_out, label.to(torch.long), weight=None)\n",
    "            loss_main = criterion[1](main_out, label)\n",
    "            loss = loss_aux + loss_main\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            tq.update(args.batch_size)\n",
    "            train_loss += loss.item()\n",
    "            tq.set_postfix(loss='%.6f' % (train_loss / (i + 1)))\n",
    "            step += 1\n",
    "            if step % 10 == 0:\n",
    "                writer.add_scalar('Train/loss_step', loss, step)\n",
    "            loss_record.append(loss.item())\n",
    "        tq.close()\n",
    "        loss_train_mean = np.mean(loss_record)\n",
    "        writer.add_scalar('Train/loss_epoch', float(loss_train_mean), epoch)\n",
    "        print('loss for train : %f' % (loss_train_mean))\n",
    "\n",
    "        if epoch % args.validation_step == 0:\n",
    "            mean_Dice, Dice1, Dice2, Dice3, acc = val(args, model, dataloader_val)\n",
    "            writer.add_scalar('Valid/Mean_val', mean_Dice, epoch)\n",
    "            writer.add_scalar('Valid/Dice1_val', Dice1, epoch)\n",
    "            writer.add_scalar('Valid/Dice2_val', Dice2, epoch)\n",
    "            writer.add_scalar('Valid/Dice3_val', Dice3, epoch)\n",
    "            writer.add_scalar('Valid/Acc_val', acc, epoch)\n",
    "            # mean_Dice=(Dice1+Dice2+Dice3)/3.0\n",
    "            is_best = mean_Dice > best_pred\n",
    "            best_pred = max(best_pred, mean_Dice)\n",
    "            print(best_pred)\n",
    "            checkpoint_dir = args.save_model_path\n",
    "            \n",
    "            if not os.path.exists(checkpoint_dir):\n",
    "                os.makedirs(checkpoint_dir)\n",
    "            checkpoint_latest = os.path.join(checkpoint_dir, 'checkpoint_latest.pth')\n",
    "            save_checkpoint({\n",
    "                'epoch': epoch + 1,\n",
    "                'state_dict': model.state_dict(),\n",
    "                'best_dice': best_pred,\n",
    "            }, best_pred, epoch, is_best, checkpoint_dir, filename=checkpoint_latest)\n",
    "\n",
    "#对测试集进行处理，变成mast\n",
    "def test(model, dataloader, args, save_path):\n",
    "    print('start test!')\n",
    "    with torch.no_grad():\n",
    "        model.eval()\n",
    "        # precision_record = []\n",
    "        tq = tqdm.tqdm(dataloader, desc='\\r')\n",
    "        tq.set_description('test')\n",
    "        comments = os.getcwd().split('\\\\')[-1]\n",
    "        for i, (data, label_path) in enumerate(tq):\n",
    "            if torch.cuda.is_available() and args.use_gpu:\n",
    "                data = data.cuda()\n",
    "                # label = label.cuda()\n",
    "            aux_pred, predict = model(data)\n",
    "            predict = torch.argmax(torch.exp(predict), dim=1)\n",
    "            pred = predict.data.cpu().numpy()\n",
    "            pred_RGB = Data.COLOR_DICT[pred.astype(np.uint8)]\n",
    "\n",
    "            for index, item in enumerate(label_path):\n",
    "                img = Image.fromarray(pred_RGB[index].squeeze().astype(np.uint8))\n",
    "                _, name = os.path.split(item)\n",
    "\n",
    "                img.save(os.path.join(save_path, name))\n",
    "                # tq.set_postfix(str=str(save_img_path))\n",
    "        tq.close()\n",
    "\n",
    "\n",
    "def main(mode='train', args=None):\n",
    "    # create dataset and dataloader\n",
    "    dataset_path = args.data\n",
    "    #训练\n",
    "    dataset_train = Data(os.path.join(dataset_path, 'train'), scale=(args.crop_width, args.crop_height), mode='train')\n",
    "    dataloader_train = DataLoader(\n",
    "        dataset_train,\n",
    "        batch_size=args.batch_size,\n",
    "        shuffle=True,\n",
    "        num_workers=args.num_workers,\n",
    "        pin_memory=True\n",
    "    )\n",
    "    #评估\n",
    "    dataset_val = Data(os.path.join(dataset_path, 'val'), scale=(args.crop_height, args.crop_width), mode='val')\n",
    "    dataloader_val = DataLoader(\n",
    "        dataset_val,\n",
    "        # this has to be 1\n",
    "        batch_size=1,\n",
    "        shuffle=True,\n",
    "        num_workers=args.num_workers,\n",
    "        pin_memory=True,\n",
    "        drop_last=True\n",
    "    )\n",
    "    #对测试集进行处理，变成mast\n",
    "    dataset_test = Data(os.path.join(dataset_path, 'test'), scale=(args.crop_height, args.crop_width), mode='test')\n",
    "    dataloader_test = DataLoader(\n",
    "        dataset_test,\n",
    "        # this has to be 1\n",
    "        batch_size=1,\n",
    "        shuffle=False,\n",
    "        num_workers=args.num_workers,\n",
    "        pin_memory=True,\n",
    "        drop_last=True\n",
    "    )\n",
    "    # build model\n",
    "    os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda\n",
    "\n",
    "    #load model\n",
    "    model_all = {'Unet': UNet()}\n",
    "    model = model_all[args.net_work]\n",
    "    print(args.net_work)\n",
    "    cudnn.benchmark = True\n",
    "    if torch.cuda.is_available() and args.use_gpu:\n",
    "        model = torch.nn.DataParallel(model).cuda()\n",
    "    optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)\n",
    "#     optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)\n",
    "    criterion_aux = nn.NLLLoss(weight=None)\n",
    "    criterion_main = Multi_DiceLoss(class_num=args.num_classes)\n",
    "    criterion = [criterion_aux, criterion_main]\n",
    "    if mode == 'train':\n",
    "        train(args, model, optimizer, criterion, dataloader_train, dataloader_val)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unet\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 0, lr 0.010000: 100%|███████████████████████████████████████████| 216/216 [01:54<00:00,  1.88it/s, loss=0.764338]\n",
      ":   0%|                                                                                       | 0/215 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss for train : 0.764338\n",
      "\n",
      "\n",
      "Start Validation!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mean_D: 0.348778, Dice1: 0.349, Dice2: 0.000, Dice3: 0.000, ACC: 0.859: 100%|████████| 215/215 [00:18<00:00, 11.78it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean_Dice: 0.34877768031172673\n",
      "Dice1: 0.3487777577420469\n",
      "Dice2: 0.0\n",
      "Dice3: 0.0\n",
      "Acc: 0.8588252374896179\n",
      "0.34877768031172673\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 1, lr 0.009699: 100%|███████████████████████████████████████████| 216/216 [01:18<00:00,  2.74it/s, loss=0.682429]\n",
      ":   1%|█                                                                              | 3/215 [00:00<00:12, 17.29it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss for train : 0.682429\n",
      "\n",
      "\n",
      "Start Validation!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mean_D: 0.551953, Dice1: 0.552, Dice2: 0.000, Dice3: 0.000, ACC: 0.856: 100%|████████| 215/215 [00:16<00:00, 12.98it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean_Dice: 0.5519527940138886\n",
      "Dice1: 0.5519528692814217\n",
      "Dice2: 0.0\n",
      "Dice3: 0.0\n",
      "Acc: 0.8559353343672135\n",
      "0.5519527940138886\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 2, lr 0.009398: 100%|███████████████████████████████████████████| 216/216 [01:20<00:00,  2.70it/s, loss=0.626740]\n",
      ":   1%|█                                                                              | 3/215 [00:00<00:12, 17.49it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss for train : 0.626740\n",
      "\n",
      "\n",
      "Start Validation!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mean_D: 0.571665, Dice1: 0.572, Dice2: 0.000, Dice3: 0.000, ACC: 0.824: 100%|████████| 215/215 [00:16<00:00, 12.96it/s]\n",
      "  0%|                                                                                          | 0/216 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean_Dice: 0.571664766739526\n",
      "Dice1: 0.5716648275627656\n",
      "Dice2: 0.0\n",
      "Dice3: 0.0\n",
      "Acc: 0.8244994711638289\n",
      "0.571664766739526\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 3, lr 0.009095: 100%|███████████████████████████████████████████| 216/216 [01:19<00:00,  2.70it/s, loss=0.612925]\n",
      ":   1%|█                                                                              | 3/215 [00:00<00:12, 17.39it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss for train : 0.612925\n",
      "\n",
      "\n",
      "Start Validation!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mean_D: 0.634164, Dice1: 0.634, Dice2: 0.000, Dice3: 0.000, ACC: 0.884: 100%|████████| 215/215 [00:16<00:00, 13.03it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean_Dice: 0.6341638724774511\n",
      "Dice1: 0.6341639585385511\n",
      "Dice2: 0.0\n",
      "Dice3: 0.0\n",
      "Acc: 0.8843668890157808\n",
      "0.6341638724774511\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 4, lr 0.008792: 100%|███████████████████████████████████████████| 216/216 [01:19<00:00,  2.71it/s, loss=0.576815]\n",
      ":   1%|█                                                                              | 3/215 [00:00<00:12, 17.59it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss for train : 0.576815\n",
      "\n",
      "\n",
      "Start Validation!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mean_D: 0.600495, Dice1: 0.600, Dice2: 0.000, Dice3: 0.000, ACC: 0.838: 100%|████████| 215/215 [00:16<00:00, 13.05it/s]\n",
      "epoch 5, lr 0.008487:   0%|                                                                    | 0/216 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean_Dice: 0.6004950925737718\n",
      "Dice1: 0.6004951567651802\n",
      "Dice2: 0.0\n",
      "Dice3: 0.0\n",
      "Acc: 0.8383517750077866\n",
      "0.6341638724774511\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 5, lr 0.008487: 100%|███████████████████████████████████████████| 216/216 [01:18<00:00,  2.74it/s, loss=0.575060]\n",
      ":   1%|█                                                                              | 3/215 [00:00<00:12, 17.19it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss for train : 0.575060\n",
      "\n",
      "\n",
      "Start Validation!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mean_D: 0.651869, Dice1: 0.652, Dice2: 0.000, Dice3: 0.000, ACC: 0.872: 100%|████████| 215/215 [00:16<00:00, 13.05it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean_Dice: 0.6518685913834561\n",
      "Dice1: 0.6518686669551313\n",
      "Dice2: 0.0\n",
      "Dice3: 0.0\n",
      "Acc: 0.8717128614254568\n",
      "0.6518685913834561\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 6, lr 0.008181: 100%|███████████████████████████████████████████| 216/216 [01:17<00:00,  2.79it/s, loss=0.549548]\n",
      ":   1%|█                                                                              | 3/215 [00:00<00:12, 17.29it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss for train : 0.549548\n",
      "\n",
      "\n",
      "Start Validation!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mean_D: 0.538658, Dice1: 0.539, Dice2: 0.000, Dice3: 0.000, ACC: 0.888: 100%|████████| 215/215 [00:16<00:00, 13.05it/s]\n",
      "epoch 7, lr 0.007873:   0%|                                                                    | 0/216 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean_Dice: 0.5386579529585025\n",
      "Dice1: 0.5386580508926386\n",
      "Dice2: 0.0\n",
      "Dice3: 0.0\n",
      "Acc: 0.8882792222539452\n",
      "0.6518685913834561\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 7, lr 0.007873: 100%|███████████████████████████████████████████| 216/216 [01:18<00:00,  2.73it/s, loss=0.534388]\n",
      ":   1%|█                                                                              | 3/215 [00:00<00:12, 17.19it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss for train : 0.534388\n",
      "\n",
      "\n",
      "Start Validation!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mean_D: 0.671477, Dice1: 0.671, Dice2: 0.000, Dice3: 0.000, ACC: 0.894: 100%|████████| 215/215 [00:16<00:00, 13.04it/s]\n",
      "  0%|                                                                                          | 0/216 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean_Dice: 0.6714765983819863\n",
      "Dice1: 0.6714766866926708\n",
      "Dice2: 0.0\n",
      "Dice3: 0.0\n",
      "Acc: 0.8937394557205149\n",
      "0.6714765983819863\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 8, lr 0.007564: 100%|███████████████████████████████████████████| 216/216 [01:18<00:00,  2.74it/s, loss=0.496241]\n",
      ":   1%|█                                                                              | 3/215 [00:00<00:12, 17.29it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss for train : 0.496241\n",
      "\n",
      "\n",
      "Start Validation!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mean_D: 0.630887, Dice1: 0.631, Dice2: 0.000, Dice3: 0.000, ACC: 0.899: 100%|████████| 215/215 [00:16<00:00, 13.07it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean_Dice: 0.6308867776506323\n",
      "Dice1: 0.6308868770998906\n",
      "Dice2: 0.0\n",
      "Dice3: 0.0\n",
      "Acc: 0.8994813431140989\n",
      "0.6714765983819863\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 9, lr 0.007254: 100%|███████████████████████████████████████████| 216/216 [01:19<00:00,  2.70it/s, loss=0.480318]\n",
      ":   1%|█                                                                              | 3/215 [00:00<00:12, 17.39it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss for train : 0.480318\n",
      "\n",
      "\n",
      "Start Validation!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mean_D: 0.675172, Dice1: 0.675, Dice2: 0.000, Dice3: 0.000, ACC: 0.905: 100%|████████| 215/215 [00:16<00:00, 13.04it/s]\n",
      "  0%|                                                                                          | 0/216 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean_Dice: 0.6751722182187267\n",
      "Dice1: 0.6751723159677978\n",
      "Dice2: 0.0\n",
      "Dice3: 0.0\n",
      "Acc: 0.9046313719892026\n",
      "0.6751722182187267\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 10, lr 0.006943: 100%|██████████████████████████████████████████| 216/216 [01:20<00:00,  2.67it/s, loss=0.488382]\n",
      ":   1%|█                                                                              | 3/215 [00:00<00:12, 17.59it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss for train : 0.488382\n",
      "\n",
      "\n",
      "Start Validation!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mean_D: 0.661809, Dice1: 0.662, Dice2: 0.000, Dice3: 0.000, ACC: 0.906: 100%|████████| 215/215 [00:16<00:00, 13.10it/s]\n",
      "epoch 11, lr 0.006629:   0%|                                                                   | 0/216 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean_Dice: 0.661808807396914\n",
      "Dice1: 0.6618089091043364\n",
      "Dice2: 0.0\n",
      "Dice3: 0.0\n",
      "Acc: 0.9061941153187293\n",
      "0.6751722182187267\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 11, lr 0.006629: 100%|██████████████████████████████████████████| 216/216 [01:17<00:00,  2.78it/s, loss=0.482556]\n",
      ":   1%|█                                                                              | 3/215 [00:00<00:12, 17.19it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss for train : 0.482556\n",
      "\n",
      "\n",
      "Start Validation!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mean_D: 0.683047, Dice1: 0.683, Dice2: 0.000, Dice3: 0.000, ACC: 0.885: 100%|████████| 215/215 [00:16<00:00, 12.97it/s]\n",
      "  0%|                                                                                          | 0/216 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean_Dice: 0.6830469019946118\n",
      "Dice1: 0.6830469816095824\n",
      "Dice2: 0.0\n",
      "Dice3: 0.0\n",
      "Acc: 0.8846033647606936\n",
      "0.6830469019946118\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 12, lr 0.006314: 100%|██████████████████████████████████████████| 216/216 [01:19<00:00,  2.73it/s, loss=0.466695]\n",
      ":   1%|█                                                                              | 3/215 [00:00<00:12, 17.49it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss for train : 0.466695\n",
      "\n",
      "\n",
      "Start Validation!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mean_D: 0.705436, Dice1: 0.705, Dice2: 0.000, Dice3: 0.000, ACC: 0.910: 100%|████████| 215/215 [00:16<00:00, 13.07it/s]\n",
      "epoch 13, lr 0.005998:   0%|                                                                   | 0/216 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean_Dice: 0.7054361257915698\n",
      "Dice1: 0.7054362236395973\n",
      "Dice2: 0.0\n",
      "Dice3: 0.0\n",
      "Acc: 0.9102788880813953\n",
      "0.7054361257915698\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 13, lr 0.005998: 100%|██████████████████████████████████████████| 216/216 [01:17<00:00,  2.79it/s, loss=0.463147]\n",
      ":   1%|█                                                                              | 3/215 [00:00<00:12, 17.59it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss for train : 0.463147\n",
      "\n",
      "\n",
      "Start Validation!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mean_D: 0.646576, Dice1: 0.647, Dice2: 0.000, Dice3: 0.000, ACC: 0.905: 100%|████████| 215/215 [00:16<00:00, 13.08it/s]\n",
      "epoch 14, lr 0.005679:   0%|                                                                   | 0/216 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean_Dice: 0.646575756456246\n",
      "Dice1: 0.6465758597050476\n",
      "Dice2: 0.0\n",
      "Dice3: 0.0\n",
      "Acc: 0.9053355270841985\n",
      "0.7054361257915698\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 14, lr 0.005679: 100%|██████████████████████████████████████████| 216/216 [01:21<00:00,  2.66it/s, loss=0.440342]\n",
      ":   1%|▋                                                                              | 2/215 [00:00<00:10, 19.66it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss for train : 0.440342\n",
      "\n",
      "\n",
      "Start Validation!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mean_D: 0.675830, Dice1: 0.676, Dice2: 0.000, Dice3: 0.000, ACC: 0.911: 100%|████████| 215/215 [00:16<00:00, 13.06it/s]\n",
      "epoch 15, lr 0.005359:   0%|                                                                   | 0/216 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean_Dice: 0.6758300078916759\n",
      "Dice1: 0.6758301120134448\n",
      "Dice2: 0.0\n",
      "Dice3: 0.0\n",
      "Acc: 0.9105748984504776\n",
      "0.7054361257915698\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 15, lr 0.005359: 100%|██████████████████████████████████████████| 216/216 [01:20<00:00,  2.68it/s, loss=0.479979]\n",
      ":   1%|█                                                                              | 3/215 [00:00<00:12, 17.39it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss for train : 0.479979\n",
      "\n",
      "\n",
      "Start Validation!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mean_D: 0.719959, Dice1: 0.720, Dice2: 0.000, Dice3: 0.000, ACC: 0.912: 100%|████████| 215/215 [00:16<00:00, 13.08it/s]\n",
      "epoch 16, lr 0.005036:   0%|                                                                   | 0/216 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean_Dice: 0.7199588819012299\n",
      "Dice1: 0.7199589783234139\n",
      "Dice2: 0.0\n",
      "Dice3: 0.0\n",
      "Acc: 0.9119015310293812\n",
      "0.7199588819012299\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 16, lr 0.005036: 100%|██████████████████████████████████████████| 216/216 [01:18<00:00,  2.75it/s, loss=0.425435]\n",
      ":   1%|█                                                                              | 3/215 [00:00<00:12, 17.45it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss for train : 0.425435\n",
      "\n",
      "\n",
      "Start Validation!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mean_D: 0.738598, Dice1: 0.739, Dice2: 0.000, Dice3: 0.000, ACC: 0.907: 100%|████████| 215/215 [00:16<00:00, 13.04it/s]\n",
      "  0%|                                                                                          | 0/216 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean_Dice: 0.7385980860762918\n",
      "Dice1: 0.7385981734563467\n",
      "Dice2: 0.0\n",
      "Dice3: 0.0\n",
      "Acc: 0.9072203575970723\n",
      "0.7385980860762918\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 17, lr 0.004711: 100%|██████████████████████████████████████████| 216/216 [01:19<00:00,  2.73it/s, loss=0.477304]\n",
      ":   1%|█                                                                              | 3/215 [00:00<00:12, 17.38it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss for train : 0.477304\n",
      "\n",
      "\n",
      "Start Validation!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mean_D: 0.708550, Dice1: 0.709, Dice2: 0.000, Dice3: 0.000, ACC: 0.915: 100%|████████| 215/215 [00:16<00:00, 13.07it/s]\n",
      "epoch 18, lr 0.004384:   0%|                                                                   | 0/216 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean_Dice: 0.708549552418633\n",
      "Dice1: 0.7085496555043501\n",
      "Dice2: 0.0\n",
      "Dice3: 0.0\n",
      "Acc: 0.9154163854858804\n",
      "0.7385980860762918\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 18, lr 0.004384: 100%|██████████████████████████████████████████| 216/216 [01:18<00:00,  2.74it/s, loss=0.442642]\n",
      ":   1%|█                                                                              | 3/215 [00:00<00:12, 17.49it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss for train : 0.442642\n",
      "\n",
      "\n",
      "Start Validation!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mean_D: 0.737721, Dice1: 0.738, Dice2: 0.000, Dice3: 0.000, ACC: 0.907: 100%|████████| 215/215 [00:16<00:00, 13.02it/s]\n",
      "epoch 19, lr 0.004054:   0%|                                                                   | 0/216 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean_Dice: 0.7377213165430894\n",
      "Dice1: 0.7377214035900675\n",
      "Dice2: 0.0\n",
      "Dice3: 0.0\n",
      "Acc: 0.9066493848629568\n",
      "0.7385980860762918\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 19, lr 0.004054: 100%|██████████████████████████████████████████| 216/216 [01:18<00:00,  2.76it/s, loss=0.420232]\n",
      ":   1%|█                                                                              | 3/215 [00:00<00:11, 17.80it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss for train : 0.420232\n",
      "\n",
      "\n",
      "Start Validation!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mean_D: 0.712406, Dice1: 0.712, Dice2: 0.000, Dice3: 0.000, ACC: 0.890: 100%|████████| 215/215 [00:16<00:00, 13.08it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean_Dice: 0.7124063616805214\n",
      "Dice1: 0.7124064401879175\n",
      "Dice2: 0.0\n",
      "Dice3: 0.0\n",
      "Acc: 0.8898898203903655\n",
      "0.7385980860762918\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 20, lr 0.003720: 100%|██████████████████████████████████████████| 216/216 [01:16<00:00,  2.83it/s, loss=0.424649]\n",
      ":   1%|█                                                                              | 3/215 [00:00<00:11, 17.69it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss for train : 0.424649\n",
      "\n",
      "\n",
      "Start Validation!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mean_D: 0.755713, Dice1: 0.756, Dice2: 0.000, Dice3: 0.000, ACC: 0.919: 100%|████████| 215/215 [00:16<00:00, 13.07it/s]\n",
      "  0%|                                                                                          | 0/216 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean_Dice: 0.7557129973977477\n",
      "Dice1: 0.7557130930275823\n",
      "Dice2: 0.0\n",
      "Dice3: 0.0\n",
      "Acc: 0.9191790710651994\n",
      "0.7557129973977477\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 21, lr 0.003384: 100%|██████████████████████████████████████████| 216/216 [01:18<00:00,  2.75it/s, loss=0.402605]\n",
      ":   1%|█                                                                              | 3/215 [00:00<00:12, 17.48it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss for train : 0.402605\n",
      "\n",
      "\n",
      "Start Validation!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mean_D: 0.751950, Dice1: 0.752, Dice2: 0.000, Dice3: 0.000, ACC: 0.915: 100%|████████| 215/215 [00:16<00:00, 13.06it/s]\n",
      "epoch 22, lr 0.003043:   0%|                                                                   | 0/216 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean_Dice: 0.7519500143510868\n",
      "Dice1: 0.7519501059474695\n",
      "Dice2: 0.0\n",
      "Dice3: 0.0\n",
      "Acc: 0.9146922773697052\n",
      "0.7557129973977477\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 22, lr 0.003043: 100%|██████████████████████████████████████████| 216/216 [01:20<00:00,  2.69it/s, loss=0.399825]\n",
      ":   1%|█                                                                              | 3/215 [00:00<00:11, 17.80it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss for train : 0.399825\n",
      "\n",
      "\n",
      "Start Validation!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mean_D: 0.757376, Dice1: 0.757, Dice2: 0.000, Dice3: 0.000, ACC: 0.917: 100%|████████| 215/215 [00:16<00:00, 13.03it/s]\n",
      "  0%|                                                                                          | 0/216 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean_Dice: 0.7573762770594059\n",
      "Dice1: 0.7573763699788829\n",
      "Dice2: 0.0\n",
      "Dice3: 0.0\n",
      "Acc: 0.9172294337884136\n",
      "0.7573762770594059\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 23, lr 0.002699: 100%|██████████████████████████████████████████| 216/216 [01:18<00:00,  2.76it/s, loss=0.375503]\n",
      ":   1%|█                                                                              | 3/215 [00:00<00:11, 17.80it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss for train : 0.375503\n",
      "\n",
      "\n",
      "Start Validation!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mean_D: 0.748440, Dice1: 0.748, Dice2: 0.000, Dice3: 0.000, ACC: 0.918: 100%|████████| 215/215 [00:16<00:00, 13.06it/s]\n",
      "epoch 24, lr 0.002349:   0%|                                                                   | 0/216 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean_Dice: 0.7484404153560257\n",
      "Dice1: 0.74844051203371\n",
      "Dice2: 0.0\n",
      "Dice3: 0.0\n",
      "Acc: 0.9183657421226121\n",
      "0.7573762770594059\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 24, lr 0.002349: 100%|██████████████████████████████████████████| 216/216 [01:19<00:00,  2.72it/s, loss=0.387534]\n",
      ":   1%|█                                                                              | 3/215 [00:00<00:12, 17.49it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss for train : 0.387534\n",
      "\n",
      "\n",
      "Start Validation!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mean_D: 0.762143, Dice1: 0.762, Dice2: 0.000, Dice3: 0.000, ACC: 0.919: 100%|████████| 215/215 [00:16<00:00, 13.05it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean_Dice: 0.7621427811629349\n",
      "Dice1: 0.7621428749908924\n",
      "Dice2: 0.0\n",
      "Dice3: 0.0\n",
      "Acc: 0.9192001596241695\n",
      "0.7621427811629349\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 25, lr 0.001994: 100%|██████████████████████████████████████████| 216/216 [01:21<00:00,  2.65it/s, loss=0.390712]\n",
      ":   1%|█                                                                              | 3/215 [00:00<00:12, 17.62it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss for train : 0.390712\n",
      "\n",
      "\n",
      "Start Validation!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mean_D: 0.717137, Dice1: 0.717, Dice2: 0.000, Dice3: 0.000, ACC: 0.919: 100%|████████| 215/215 [00:16<00:00, 13.08it/s]\n",
      "epoch 26, lr 0.001631:   0%|                                                                   | 0/216 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean_Dice: 0.7171368366188163\n",
      "Dice1: 0.7171369426569926\n",
      "Dice2: 0.0\n",
      "Dice3: 0.0\n",
      "Acc: 0.9193581615967608\n",
      "0.7621427811629349\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 26, lr 0.001631: 100%|██████████████████████████████████████████| 216/216 [01:17<00:00,  2.78it/s, loss=0.373415]\n",
      ":   1%|█                                                                              | 3/215 [00:00<00:12, 17.59it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss for train : 0.373415\n",
      "\n",
      "\n",
      "Start Validation!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mean_D: 0.747306, Dice1: 0.747, Dice2: 0.000, Dice3: 0.000, ACC: 0.924: 100%|████████| 215/215 [00:16<00:00, 13.07it/s]\n",
      "epoch 27, lr 0.001259:   0%|                                                                   | 0/216 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean_Dice: 0.7473055866408455\n",
      "Dice1: 0.7473056901816809\n",
      "Dice2: 0.0\n",
      "Dice3: 0.0\n",
      "Acc: 0.9235340612671304\n",
      "0.7621427811629349\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 27, lr 0.001259: 100%|██████████████████████████████████████████| 216/216 [01:19<00:00,  2.72it/s, loss=0.349769]\n",
      ":   1%|█                                                                              | 3/215 [00:00<00:12, 17.25it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss for train : 0.349769\n",
      "\n",
      "\n",
      "Start Validation!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mean_D: 0.765491, Dice1: 0.765, Dice2: 0.000, Dice3: 0.000, ACC: 0.925: 100%|████████| 215/215 [00:16<00:00, 13.04it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean_Dice: 0.765491000147951\n",
      "Dice1: 0.7654910996575922\n",
      "Dice2: 0.0\n",
      "Dice3: 0.0\n",
      "Acc: 0.924598019622093\n",
      "0.765491000147951\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 28, lr 0.000874: 100%|██████████████████████████████████████████| 216/216 [01:16<00:00,  2.84it/s, loss=0.348833]\n",
      ":   1%|▋                                                                              | 2/215 [00:00<00:10, 19.66it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss for train : 0.348833\n",
      "\n",
      "\n",
      "Start Validation!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mean_D: 0.761280, Dice1: 0.761, Dice2: 0.000, Dice3: 0.000, ACC: 0.926: 100%|████████| 215/215 [00:16<00:00, 13.07it/s]\n",
      "epoch 29, lr 0.000468:   0%|                                                                   | 0/216 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean_Dice: 0.7612804557835049\n",
      "Dice1: 0.7612805586072566\n",
      "Dice2: 0.0\n",
      "Dice3: 0.0\n",
      "Acc: 0.9260749081836587\n",
      "0.765491000147951\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 29, lr 0.000468: 100%|██████████████████████████████████████████| 216/216 [01:18<00:00,  2.74it/s, loss=0.361801]\n",
      ":   1%|█                                                                              | 3/215 [00:00<00:12, 17.64it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss for train : 0.361801\n",
      "\n",
      "\n",
      "Start Validation!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Mean_D: 0.761803, Dice1: 0.762, Dice2: 0.000, Dice3: 0.000, ACC: 0.917: 100%|████████| 215/215 [00:16<00:00, 13.06it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mean_Dice: 0.761802506479558\n",
      "Dice1: 0.7618025981752353\n",
      "Dice2: 0.0\n",
      "Dice3: 0.0\n",
      "Acc: 0.9172352331421304\n",
      "0.765491000147951\n"
     ]
    }
   ],
   "source": [
    "if __name__ == '__main__':\n",
    "    seed = 1234\n",
    "    torch.manual_seed(seed)  # 固定初始化\n",
    "    torch.cuda.manual_seed_all(seed)\n",
    "    args = DefaultConfig()\n",
    "    modes = 'train'\n",
    "    if modes == 'train':\n",
    "        main(mode='train', args=args)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 预测"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "test:   0%|                                                                                    | 0/107 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "start test!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "test: 100%|██████████████████████████████████████████████████████████████████████████| 107/107 [00:10<00:00, 10.52it/s]\n"
     ]
    }
   ],
   "source": [
    "if __name__ == '__main__':\n",
    "    save_path = r'testjie'\n",
    "    model_path = r'checkpoints/checkpoint_latest.pth'\n",
    "    dataset_test = Data(os.path.join(DefaultConfig.data, 'test'), scale=(DefaultConfig.crop_width,\n",
    "                                                                         DefaultConfig.crop_height), mode='test')\n",
    "    args = DefaultConfig()\n",
    "    dataloader_test = DataLoader(\n",
    "        dataset_test,\n",
    "        batch_size=1,\n",
    "        shuffle=False,\n",
    "        num_workers=0,\n",
    "        pin_memory=True,\n",
    "        drop_last=False\n",
    "    )\n",
    "    os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda\n",
    "    model_all = {'Unet':UNet()}\n",
    "    model = model_all[args.net_work]\n",
    "    cudnn.benchmark = True\n",
    "    if torch.cuda.is_available() and args.use_gpu:\n",
    "        model = torch.nn.DataParallel(model).cuda()\n",
    "    checkpoint = torch.load(model_path)\n",
    "    model.load_state_dict(checkpoint['state_dict'])\n",
    "    test(model, dataloader_test, args, save_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "其他的"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "4905652b14e4b7eb92899b78ac499a22c488804455b27940a322fd82aaf71031"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
